Routed per-subset state dict gets into NIXL CUDA destinations - #252

Root issue: torchtitan RL generators prefetch each policy into a full pinned CPU copy of the weights, and this copy OOMs GB300 generator hosts, which have little CPU memory. The companion torchtitan PR reads the StorageVolume snapshot straight into the generator's GPU parameters, one FSDP group at a time. Two pieces were missing in torchstore:

• Routed per-subset gets: RoutingClient.get_layouts planned only the first state dict it saw, so a later get_state_dict for another subset returned the first subset's keys in fresh tensors and none of the requested ones.
• Safe NIXL CUDA destinations: FSDP frees and reallocates sharded param storage in place with resize_, so a cached registration ends up over freed memory and reads stale bytes. With expandable segments, UCX also registers only the 20 MiB cuMem chunk holding a buffer's start. A buffer that fits within one chunk but crosses a chunk boundary then fails registration (ibv_reg_dmabuf_mr: Invalid argument on GB200) or, past the registered part, RDMA with a local access violation that hangs the transfer. UCX's registration cache also keeps deregistered memory registered, so CUDA memory torch frees after a get stays pinned and the device cannot reuse it (measured: freed destinations missing from cudaMemGetInfo until the process exits).

Fix:

• RoutingClient.get_layouts plans every key the first time it is passed, against only that subset's publisher entries, and fetches publisher layouts once per state dict. get_state_dict(strict=False) fetches only the passed subset; keys it skips as unpublished stay unplanned, so a later strict get still rejects them.
• NIXL CUDA get destinations hold a registration of their whole storage only while the get is in flight. Concurrent buffers of one get share it. CPU destinations stay cached. A storage is never registered twice: two registrations of the same range make peers fail to load the agent's metadata (NIXL_ERR_NOT_ALLOWED).
• Cancelling a CUDA get waits, up to TORCHSTORE_NIXL_TIMEOUT_S, for the storage volume to answer before deregistering its destinations, since the volume may still be writing into them. After a timeout or a second cancel, the destinations stay registered and referenced until the volume answers, and getting into the same storage meanwhile raises. This uses a new TransportBuffer._call_volume_get hook.
• The NIXL UCX backend is created with engine_config CUDA_COPY_REG_WHOLE_ALLOC=off,RCACHE_ENABLE=n, scoped to torchstore's agent: every registration covers its whole range, and deregistering releases the memory. Exported UCX_CUDA_COPY_REG_WHOLE_ALLOC/UCX_RCACHE_ENABLE still override these.

BC: no API change. Routed non-strict gets return the passed keys rather than those of the first get. Routed strict gets still require the passed dict to match every key planned so far.

Companion torchtitan PR (lands after this one): pytorch/torchtitan#5652. Duplicate trainer copy follow-up: pytorch/torchtitan#5654

Test plan:

• pytest tests/routing tests/test_nixl_transport.py tests/test_nixl_transport_cuda.py tests/test_state_dict.py plus tests/test_state_dict_utils.py, rebased on main, on GB200 with nixl 1.5.0: 135 passed, 1 failed. The failure, test_state_dict_transfer_dtype[MonarchRDMA-cuda] (RDMA completion error status=12 in the put path), also fails 5/5 on main. Cancellation tests cover waiting for the volume's answer, release after a timeout or second cancel, and refusing to reuse a storage the volume may still write into. The CUDA tests use real NIXL and cover a destination spanning cuMem chunks with expandable segments (fails registration without CUDA_COPY_REG_WHOLE_ALLOC=off) and a freed destination returning to the device (fails without RCACHE_ENABLE=n).
• Two-process check (CPU volume, GPU client, IB): NIXL allocates no GPU memory of its own; device usage matches a plain copy_ baseline, does not grow over 20 gets, and freed destinations return to the device. NIXL segfaults at interpreter exit on this host for any process that creates an agent, so judge by the pytest summary.
• End to end, torchtitan RL fetch_model_state_dict_to_gpu (Qwen3.5, GB200): rc=0, logprob_diff 0.010-0.014, matching the staged path.

Limitations:

• CUDA destinations register and deregister on every get, including callers whose storage never moves. Without the UCX registration cache, GPU register+deregister measured 0.9 ms (64 MiB), 2.8 ms (1 GiB), 9.1 ms (4 GiB) on GB200, about 1.5 s per sync for 1.5k params plus ~3 ms/GiB; not measured at scale. Caching by (data_ptr, nbytes) cannot detect FSDP reallocation at the same address.
• tests/test_nixl_transport_cuda.py skips in CI (nixl is not installed); it was run manually.
• RCACHE_ENABLE=n applies to all memory on both agents, so uncached registrations (per-get CPU destinations, staging buffers, non-contiguous copies) pay a full register/deregister each time. Long-lived registrations are cached by torchstore.
• Each CUDA get changes the client's NIXL metadata, so the volume reloads it per get, as it already did for new CPU destinations. Reloading removes the remote agent, so concurrent gets from one client to one volume can hang an in-flight transfer (pre-existing; now hit by any concurrent CUDA gets). torchtitan issues one blocking get at a time.
• A get abandoned after a timeout keeps its destinations registered until the volume answers. Holding the tensors does not stop FSDP resize_ from freeing that memory, and gets refuse the storage only if it comes back at the same address and size.
• A first routed get whose subset has no published keys still fails (Mapping is missing); later gets now recover.

Lines changed: impl +218/-14, tests +460/-2.